// Copyright (c) 2021 Presto Labs Pte. Ltd.
// Author: xguo

#include "coin2/exchange/dydx_futures/api_util/auth_key.h"

#include <bitset>
#include <cmath>
#include <string>
#include <tuple>
#include <unordered_map>

#include <nlohmann/json.hpp>
#include <openssl/bn.h>
#include <openssl/ec.h>

// #include "coin2/exchange/base/api_base/asio_http_driver.h"
#include "crypto/ecdsa.h"
#include "crypto/openssl_helper.h"
#include "crypto/pedersen_hash.h"
#include "starkware/crypto/elliptic_curve_constants.h"

#define VERIFY_SIGNATURE false

namespace {
using namespace starkware;
constexpr auto vt_upper_bound = 0x800000000000000000000000000000000000000000000000000000000000000_Z;
}  // namespace

// From official python code.
// https://github.com/dydxprotocol/dydx-v3-python/blob/master/dydx3/constants.py
//
std::unordered_map<std::string, std::string> SYNTHETIC_ASSET_MAP = {
    {"BTC-USD", "BTC"},     {"ETH-USD", "ETH"},     {"LINK-USD", "LINK"}, {"AAVE-USD", "AAVE"},
    {"UNI-USD", "UNI"},     {"SUSHI-USD", "SUSHI"}, {"SOL-USD", "SOL"},   {"YFI-USD", "YFI"},
    {"1INCH-USD", "1INCH"}, {"AVAX-USD", "AVAX"},   {"SNX-USD", "SNX"},   {"CRV-USD", "CRV"},
    {"UMA-USD", "UMA"},     {"DOT-USD", "DOT"},     {"DOGE-USD", "DOGE"}, {"MATIC-USD", "MATIC"},
    {"MKR-USD", "MKR"},     {"FIL-USD", "FIL"},     {"ADA-USD", "ADA"},   {"ATOM-USD", "ATOM"},
    {"COMP-USD", "COMP"},   {"BCH-USD", "BCH"},     {"LTC-USD", "LTC"},   {"EOS-USD", "EOS"},
    {"ALGO-USD", "ALGO"},   {"ZRX-USD", "ZRX"},     {"XMR-USD", "XMR"},   {"ZEC-USD", "ZEC"},
    {"ENJ-USD", "ENJ"},     {"ETC-USD", "ETC"},     {"XLM-USD", "XLM"},   {"TRX-USD", "TRX"},
    {"XTZ-USD", "XTZ"},     {"ICP-USD", "ICP"},     {"RUNE-USD", "RUNE"}, {"LUNA-USD", "LUNA"},
    {"NEAR-USD", "NEAR"},   {"CELO-USD", "CELO"}};

std::unordered_map<std::string, const starkware::BigInt<4>> SYNTHETIC_ASSET_ID_MAP = {
    {"BTC", 0x4254432d3130000000000000000000_Z},   {"ETH", 0x4554482d3900000000000000000000_Z},
    {"LINK", 0x4c494e4b2d37000000000000000000_Z},  {"AAVE", 0x414156452d38000000000000000000_Z},
    {"UNI", 0x554e492d3700000000000000000000_Z},   {"SUSHI", 0x53555348492d370000000000000000_Z},
    {"SOL", 0x534f4c2d3700000000000000000000_Z},   {"YFI", 0x5946492d3130000000000000000000_Z},
    {"1INCH", 0x31494e43482d370000000000000000_Z}, {"AVAX", 0x415641582d37000000000000000000_Z},
    {"SNX", 0x534e582d3700000000000000000000_Z},   {"CRV", 0x4352562d3600000000000000000000_Z},
    {"UMA", 0x554d412d3700000000000000000000_Z},   {"DOT", 0x444f542d3700000000000000000000_Z},
    {"DOGE", 0x444f47452d35000000000000000000_Z},  {"MATIC", 0x4d415449432d360000000000000000_Z},
    {"MKR", 0x4d4b522d3900000000000000000000_Z},   {"FIL", 0x46494c2d3700000000000000000000_Z},
    {"ADA", 0x4144412d3600000000000000000000_Z},   {"ATOM", 0x41544f4d2d37000000000000000000_Z},
    {"COMP", 0x434f4d502d38000000000000000000_Z},  {"BCH", 0x4243482d3800000000000000000000_Z},
    {"LTC", 0x4c54432d3800000000000000000000_Z},   {"EOS", 0x454f532d3600000000000000000000_Z},
    {"ALGO", 0x414c474f2d36000000000000000000_Z},  {"ZRX", 0x5a52582d3600000000000000000000_Z},
    {"XMR", 0x584d522d3800000000000000000000_Z},   {"ZEC", 0x5a45432d3800000000000000000000_Z},
    {"ENJ", 0x454e4a2d3600000000000000000000_Z},   {"ETC", 0x4554432d3700000000000000000000_Z},
    {"XLM", 0x584c4d2d3500000000000000000000_Z},   {"TRX", 0x5452582d3400000000000000000000_Z},
    {"XTZ", 0x58545a2d3600000000000000000000_Z},   {"ICP", 0x4943502d3700000000000000000000_Z},
    {"RUNE", 0x52554e452d36000000000000000000_Z},  {"LUNA", 0x4c554e412d36000000000000000000_Z},
    {"NEAR", 0x4e4541522d36000000000000000000_Z},  {"CELO", 0x43454c4f2d36000000000000000000_Z}};

std::unordered_map<std::string, double> ASSET_RESOLUTION = {
    {"USDC", 1e6},  {"BTC", 1e10}, {"ETH", 1e9},  {"LINK", 1e7},    {"AAVE", 1e8},  {"UNI", 1e7},
    {"SUSHI", 1e7}, {"SOL", 1e7},  {"YFI", 1e10}, {"ONEINCH", 1e7}, {"AVAX", 1e7},  {"SNX", 1e7},
    {"CRV", 1e6},   {"UMA", 1e7},  {"DOT", 1e7},  {"DOGE", 1e5},    {"MATIC", 1e6}, {"MKR", 1e9},
    {"FIL", 1e7},   {"ADA", 1e6},  {"ATOM", 1e7}, {"COMP", 1e8},    {"BCH", 1e8},   {"LTC", 1e8},
    {"EOS", 1e6},   {"ALGO", 1e6}, {"ZRX", 1e6},  {"XMR", 1e8},     {"ZEC", 1e8},   {"ENJ", 1e6},
    {"ETC", 1e7},   {"XLM", 1e5},  {"TRX", 1e4},  {"XTZ", 1e6},     {"ICP", 1e7},   {"RUNE", 1e6},
    {"LUNA", 1e6},  {"NEAR", 1e6}, {"CELO", 1e6}};

constexpr ValueType MAINNET_COLLATERAL =
    0x2893294412a4c8f915f75892b395ebbf6859ec246ec365c3b1f56f47c3a0a5d_Z;

constexpr int ORDER_PREFIX = 3;

uint64_t NonceFromClientId(const std::string& data) {
  std::string digest = ApiSignatureUtil::GetHashByRawSHA256(data);

  char tmp[4];
  tmp[0] = digest[31];
  tmp[1] = digest[30];
  tmp[2] = digest[29];
  tmp[3] = digest[28];
  return *(reinterpret_cast<uint32_t*>(tmp));
}

StarkwareOrder::StarkwareOrder(const OrderSpec& order_spec) {
  const auto& asset = SYNTHETIC_ASSET_MAP.at(order_spec.product);

  asset_id_synthetic = SYNTHETIC_ASSET_ID_MAP.at(asset);
  asset_id_collateral = MAINNET_COLLATERAL;
  asset_id_fee = MAINNET_COLLATERAL;

  if (order_spec.order_side == coin::proto::OrderSide::BUY_ORDER) {
    is_buy_synthetic = true;
  } else if (order_spec.order_side == coin::proto::OrderSide::SELL_ORDER) {
    is_buy_synthetic = false;
  } else {
    CHECK(false);
  }

  // CHECK decimal calculation.
  if (is_buy_synthetic) {
    quantums_amount_synthetic =
        static_cast<uint64_t>(std::round(order_spec.qty * ASSET_RESOLUTION[asset]));
    quantums_amount_collateral = static_cast<uint64_t>(
        std::round(order_spec.price * order_spec.qty * ASSET_RESOLUTION["USDC"]));
  } else {
    quantums_amount_synthetic =
        static_cast<uint64_t>(std::round(order_spec.qty * ASSET_RESOLUTION[asset]));
    quantums_amount_collateral = static_cast<uint64_t>(
        std::round(order_spec.price * order_spec.qty * ASSET_RESOLUTION["USDC"]));
  }

  // round down
  quantums_amount_fee = static_cast<uint64_t>(std::ceil(0.0015 * quantums_amount_collateral));

  expiration_epoch_hours = std::ceil((order_spec.order_expiration_ts / 1e9) / 3600) + 7 * 24;
  nonce = NonceFromClientId(std::to_string(order_spec.proc_order_id));

  ////////////////////////////////////////////////////////////////////////////////////////
  json = nlohmann::json::object();
  json["client_id"] = std::to_string(order_spec.proc_order_id);
  json["market"] = order_spec.product;
  if (order_spec.order_side == coin::proto::BUY_ORDER) {
    json["side"] = "BUY";
  } else if (order_spec.order_side == coin::proto::SELL_ORDER) {
    json["side"] = "SELL";
  } else {
    CHECK(false) << (int)order_spec.order_side;
  }

  json["human_price"] = std::to_string(order_spec.price);
  json["human_size"] = std::to_string(order_spec.qty);
  json["limit_fee"] = "0.0015";
  json["expiration_epoch_seconds"] = (order_spec.order_expiration_ts / 1e9);
  ////////////////////////////////////////////////////////////////////////////////////////
}

int StarkwareOrder::CalculateHash(const EC_GROUP* group, BIGNUM* res) {
  ValueType vt_asset_id_sell = ValueType::Zero();
  ValueType vt_asset_id_buy = ValueType::Zero();
  uint64_t quantums_amount_sell;
  uint64_t quantums_amount_buy;

  if (is_buy_synthetic) {
    vt_asset_id_sell = asset_id_collateral;
    vt_asset_id_buy = asset_id_synthetic;
    quantums_amount_sell = quantums_amount_collateral;
    quantums_amount_buy = quantums_amount_synthetic;
  } else {
    vt_asset_id_sell = asset_id_synthetic;
    vt_asset_id_buy = asset_id_collateral;
    quantums_amount_sell = quantums_amount_synthetic;
    quantums_amount_buy = quantums_amount_collateral;
  }

  // Get Pedersen Hash constants
  EC_POINT* shift_point = EC_POINT_new(group);
  std::vector<EC_POINT*> points;
  const size_t n_element_bits = 252;  // we only need 2*n_element_bits points for pedersen hash
  for (size_t j = 0; j < 2 * n_element_bits; j++) {
    points.push_back(EC_POINT_new(group));
  }
  EcConstants2EC_POINT(group, shift_point, points);

  // get part 1
  ValueType vt_part_1 = ValueType::Zero();
  auto* ptr = reinterpret_cast<uint32_t*>(&vt_part_1);
  ptr[0] = static_cast<uint32_t>(nonce);
  (reinterpret_cast<uint64_t*>(&ptr[1]))[0] = quantums_amount_fee;
  (reinterpret_cast<uint64_t*>(&ptr[3]))[0] = quantums_amount_buy;
  (reinterpret_cast<uint64_t*>(&ptr[5]))[0] = quantums_amount_sell;
  // DLOG(INFO) << "nonce: " << nonce;
  // DLOG(INFO) << "fee: " << quantums_amount_fee;
  // DLOG(INFO) << "buy: " << quantums_amount_buy;
  // DLOG(INFO) << "sell: " << quantums_amount_sell;

  // get part 2
  ValueType vt_part_2 = ValueType::Zero();
  ptr = reinterpret_cast<uint32_t*>(&vt_part_2);
  ptr[0] = static_cast<uint32_t>(expiration_epoch_hours);
  (reinterpret_cast<uint64_t*>(&ptr[1]))[0] = position_id;
  (reinterpret_cast<uint64_t*>(&ptr[3]))[0] = position_id;
  (reinterpret_cast<uint64_t*>(&ptr[5]))[0] = position_id;
  ptr[7] = static_cast<uint32_t>(ORDER_PREFIX);
  reinterpret_cast<std::bitset<256>*>(ptr)->operator<<=(17);
  // DLOG(INFO) << "part 1: " << vt_part_1;
  // DLOG(INFO) << "part 2: " << vt_part_2;

  BIGNUM* bn_asset_id_fee = BN_new();
  BIGNUM* bn_asset_id_sell = BN_new();
  BIGNUM* bn_asset_id_buy = BN_new();
  BIGNUM* bn_part_1 = BN_new();
  BIGNUM* bn_part_2 = BN_new();
  VT2BN(asset_id_fee, bn_asset_id_fee);
  VT2BN(vt_asset_id_sell, bn_asset_id_sell);
  VT2BN(vt_asset_id_buy, bn_asset_id_buy);
  VT2BN(vt_part_1, bn_part_1);
  VT2BN(vt_part_2, bn_part_2);

  BIGNUM* bn_assets_hash = BN_new();
  BIGNUM* bn_tmp = BN_new();
  PedersenHash0(group, bn_asset_id_sell, bn_asset_id_buy, shift_point, points, bn_tmp);
  PedersenHash0(group, bn_tmp, bn_asset_id_fee, shift_point, points, bn_assets_hash);
  PedersenHash0(group, bn_assets_hash, bn_part_1, shift_point, points, bn_tmp);
  PedersenHash0(group, bn_tmp, bn_part_2, shift_point, points, res);

  // DLOG(INFO) << "assets_id_sell: " << vt_asset_id_sell;
  // DLOG(INFO) << "assets_id_buy: " << vt_asset_id_buy;
  // DLOG(INFO) << "bn_assets_hash: " << BN2VT(bn_assets_hash);
  // DLOG(INFO) << "total hash: " << BN2VT(res);

  BN_free(bn_part_1);
  BN_free(bn_part_2);
  BN_free(bn_asset_id_fee);
  BN_free(bn_asset_id_sell);
  BN_free(bn_asset_id_buy);
  BN_free(bn_assets_hash);
  BN_free(bn_tmp);
  EC_POINT_free(shift_point);
  for (size_t j = 0; j < 2 * n_element_bits; j++) {
    EC_POINT_free(points[j]);
  }
  return 1;
}

std::string StarkwareOrder::Sign(const ValueType& privkey) {
  const auto& vt_curve_order = GetEcConstants().k_order;
  static_assert(vt_upper_bound <= PrimeFieldElement::kModulus);
  CHECK_LE(vt_upper_bound, vt_curve_order) << "Unexpected curve size.";
  CHECK_GT(vt_curve_order.NumLeadingZeros(), 0) << "Implementation assumes smaller curve.";

  // Convert Constants to BIGNUM
  BIGNUM* bn_modulo = BN_new();
  BIGNUM* bn_privkey = BN_new();
  BIGNUM* bn_alpha = BN_new();
  BIGNUM* bn_beta = BN_new();
  BIGNUM* bn_curve_order = BN_new();
  BIGNUM* bn_upper_bound = BN_new();
  VT2BN(PrimeFieldElement::kModulus, bn_modulo);
  VT2BN(privkey, bn_privkey);
  VT2BN(GetEcConstants().k_alpha.ToStandardForm(), bn_alpha);
  VT2BN(GetEcConstants().k_beta.ToStandardForm(), bn_beta);
  VT2BN(vt_curve_order, bn_curve_order);
  VT2BN(vt_upper_bound, bn_upper_bound);
  BN_CTX* bn_ctx = BN_CTX_new();
  EC_GROUP* group = EC_GROUP_new_curve_GFp(bn_modulo, bn_alpha, bn_beta, bn_ctx);  // Elliptic Curve

  // Public Key Point and Generator Point for ECDSA
  BIGNUM* x = BN_new();
  VT2BN(pub_key, x);
  EC_POINT* generator = EC_POINT_new(group);
  EC_POINT* public_key = EC_POINT_new(group);
  EcPoint2EC_POINT(group, GetEcConstants().k_points[1], generator);
  GetPointFromX(group, x, bn_alpha, bn_beta, bn_modulo, public_key);
  ASSERT(
      EC_POINT_is_on_curve(group, public_key, bn_ctx),
      "Given public key (" + pub_key.ToString() +
          ") does not correspond to a valid point on the elliptic curve");

  // Calculate hash
  BIGNUM* z = BN_new();
  CalculateHash(group, z);
  // DLOG(INFO) << "hash: " << BN2VT(z);
  ASSERT(!BN_is_zero(z), "Message cannot be zero.");
  ASSERT(BN_cmp(z, bn_upper_bound) == -1, "z is too big.");

  std::string res = SignEcdsa0(
      group,
      public_key,
      generator,
      bn_privkey,
      z,
      bn_curve_order,
      bn_upper_bound,
      VERIFY_SIGNATURE);  // Sign and Verify

  BN_free(bn_privkey);
  BN_free(bn_alpha);
  BN_free(bn_beta);
  BN_free(bn_upper_bound);
  BN_free(bn_curve_order);
  BN_free(bn_modulo);
  BN_free(z);
  BN_free(x);

  EC_POINT_free(public_key);
  EC_POINT_free(generator);
  EC_GROUP_free(group);
  BN_CTX_free(bn_ctx);
  return res;
}
